#!/bin/bash
##Filename:     iptables.sh
##Date:         2024-06-28-01
##Description:  这是一个对iptables做策略的脚本
##Author：		linux1106

#!/bin/bash

# 设置语言环境为UTF-8
export LANG=zh_CN.UTF-8
export LANGUAGE=zh_CN.UTF-8
export LC_ALL=zh_CN.UTF-8

# 默认配置参数
DEFAULT_LOG_FILE="/var/log/iptables_script.log"
DEFAULT_MAX_LOG_SIZE=10485760  # 10MB
DEFAULT_MAX_LOG_BACKUPS=5      # 最大日志备份数量

# 读取配置文件
CONFIG_FILE="/etc/iptables_script.conf"
if [ -f "$CONFIG_FILE" ]; then
    source "$CONFIG_FILE"
else
    LOG_FILE="$DEFAULT_LOG_FILE"
    MAX_LOG_SIZE="$DEFAULT_MAX_LOG_SIZE"
    MAX_LOG_BACKUPS="$DEFAULT_MAX_LOG_BACKUPS"
fi

# 检查是否以root用户权限运行脚本
if [ "$EUID" -ne 0 ]; then
    echo "请以root用户权限运行此脚本。"
    exit 1
fi

# 检查iptables命令是否存在
if ! command -v iptables &> /dev/null; then
    echo "找不到iptables命令，请确保已安装iptables。"
    exit 1
fi

# 检查dialog是否安装，如果未安装则尝试安装
if ! command -v dialog &> /dev/null; then
    echo "dialog未安装，正在尝试安装..."
    if command -v yum &> /dev/null; then
        yum install -y dialog
    elif command -v apt-get &> /dev/null; then
        apt-get update
        apt-get install -y dialog
    else
        echo "无法自动安装dialog，请手动安装dialog后再运行此脚本。"
        exit 1
    fi

    # 检查是否成功安装
    if ! command -v dialog &> /dev/null; then
        echo "dialog安装失败，请手动安装dialog后再运行此脚本。"
        exit 1
    fi
fi

# 通用方法获取文件大小
function file_size() {
    if [ -f "$1" ];then
        wc -c < "$1"
    else
        echo 0
    fi
}

# 日志轮转和备份管理
function rotate_log() {
    local log_size=$(file_size "$LOG_FILE")
    if [ "$log_size" -ge "$MAX_LOG_SIZE" ];then
        for (( i=$MAX_LOG_BACKUPS; i>0; i-- ));do
            if [ -f "$LOG_FILE.$i" ];then
                if [ "$i" -eq "$MAX_LOG_BACKUPS" ];then
                    rm -f "$LOG_FILE.$i"
                else
                    mv "$LOG_FILE.$i" "$LOG_FILE.$((i+1))"
                fi
            fi
        done
        mv "$LOG_FILE" "$LOG_FILE.1"
    fi
}

# 记录日志
function log() {
    rotate_log
    local log_level=$1
    shift
    local log_message="$*"
    echo "$(date '+%Y-%m-%d %H:%M:%S') [$log_level] $log_message" >> "$LOG_FILE"
}

# 验证IPv4地址或地址段
function validate_ipv4() {
    local ip=$1
    local base_ip=$(echo $ip | cut -d '/' -f 1)
    if [[ $base_ip =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}$ ]]; then
        for i in {1..4};do
            if [ "$(echo "$base_ip" | cut -d. -f$i)" -gt 255 ];then
                dialog --msgbox "IPv4地址无效。" 6 40
                return 1
            fi
        done
        return 0
    elif [[ $ip =~ ^([0-9]{1,3}\.){3}[0-9]{1,3}/([0-9]|[1-2][0-9]|3[0-2])$ ]];then
        return 0
    else
        dialog --msgbox "IPv4地址或地址段无效。" 6 40
        return 1
    fi
}

# 验证IPv6地址或地址段
function validate_ipv6() {
    local ip=$1
    if [[ $ip =~ ^([0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}$ ]] || [[ $ip =~ ^([0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}/([1-9]|[1-9][0-9]|1[0-1][0-9]|12[0-8])$ ]];then
        return 0
    else
        dialog --msgbox "IPv6地址或地址段无效。" 6 40
        return 1
    fi
}

# 验证IP地址或地址段 (IPv4 和 IPv6)
function validate_ip() {
    local ip=$1
    validate_ipv4 "$ip" || validate_ipv6 "$ip"
}

# 验证单个端口号
function validate_single_port() {
    local port=$1
    if [[ "$port" =~ ^[1-9][0-9]{0,4}$ ]] && [[ "$port" -ge 1 ]] && [[ "$port" -le 65535 ]];then
        return 0
    else
        dialog --msgbox "端口号范围无效。" 6 40
        return 1
    fi
}

# 验证端口范围
function validate_port_range() {
    local port_range=$1
    local IFS=-
    local -a ports
    read -r -a ports <<< "$port_range"
    if [[ ${#ports[@]} -eq 2 ]] && validate_single_port "${ports[0]}" && validate_single_port "${ports[1]}" && [[ "${ports[0]}" -lt "${ports[1]}" ]];then
        return 0
    else
        dialog --msgbox "端口范围格式无效。" 6 40
        return 1
    fi
}

# 验证端口号或端口范围
function validate_port() {
    local port=$1
    if [ -z "$port" ];then
        return 0
    elif validate_single_port "$port" || validate_port_range "$port";then
        return 0
    else
        return 1
    fi
}

# 验证协议类型
function validate_protocol() {
    local protocol=$1
    if [[ "$protocol" == "tcp" || "$protocol" == "udp" || "$protocol" == "icmp" || "$protocol" == "all" ]];then
        return 0
    else
        dialog --msgbox "协议无效，请输入tcp、udp、icmp或all。" 6 40
        return 1
    fi
}

# 验证操作类型
function validate_action() {
    local action=$1
    if [[ "$action" == "ACCEPT" || "$action" == "DROP" || "$action" == "REJECT" ]];then
        return 0
    else
        dialog --msgbox "操作无效，请输入ACCEPT、DROP或REJECT。" 6 40
        return 1
    fi
}

# 获取用户输入，并验证输入
function get_input() {
    local prompt=$1
    local validation_func=$2
    local default_value=$3
    local input

    while true;do
        input=$(dialog --inputbox "$prompt" 8 40 "$default_value" 3>&1 1>&2 2>&3 3>&-)
        local exit_status=$?
        if [ $exit_status -eq 1 ]; then
            dialog --msgbox "已取消操作或输入无效。" 6 40
            return 1
        fi
        if [ -z "$input" ] && [ -n "$default_value" ];then
            input=$default_value
        fi
        if [ -z "$input" ];then
            dialog --msgbox "输入不能为空，请重新输入。" 6 40
        elif ! $validation_func "$input";then
            dialog --msgbox "输入格式无效，请重试。" 6 40
        else
            echo "$input"
            return 0
        fi
    done
}

# 检查规则是否已存在
function check_rule_exists() {
    local rule=$1
    if iptables -C $rule &> /dev/null;then
        return 0
    else
        return 1
    fi
}

# 执行iptables命令并处理错误
function execute_iptables() {
    local rule=$1
    if check_rule_exists "$rule";then
        log "INFO" "iptables规则已存在: $rule"
        dialog --msgbox "iptables规则已存在。" 6 40
    else
        if ! iptables $rule 2>>"$LOG_FILE";then
            local error_code=$?
            log "ERROR" "添加iptables规则失败 (错误码: $error_code): $rule"
            dialog --msgbox "添加iptables规则失败 (错误码: $error_code)，请检查输入并重试。" 6 40
            return 1
        fi
        log "INFO" "iptables规则已成功添加: $rule"
        dialog --msgbox "iptables规则已成功添加。" 6 40
        return 0
    fi
}

# 添加iptables规则
function add_iptables_rule() {
    local src_ip=$(get_input "请输入源IP地址或地址段 (如192.168.1.0/24):" validate_ip)
    if [ $? -ne 0 ]; then
        return
    fi

    local dst_port=$(get_input "请输入目标端口或端口范围（可选，按Enter跳过）:" validate_port)
    if [ $? -ne 0 ]; then
        return
    fi

    local protocol=$(get_input "请输入协议 (tcp/udp/icmp/all):" validate_protocol)
    if [ $? -ne 0 ]; then
        return
    fi

    local action=$(get_input "请输入操作 (ACCEPT/DROP/REJECT):" validate_action)
    if [ $? -ne 0 ]; then
        return
    fi

    local rule="-A INPUT -s $src_ip"
    if [ -n "$dst_port" ]; then
        rule+=" -p $protocol --dport $dst_port"
    fi
    rule+=" -j $action"

    execute_iptables "$rule"
}

# 删除iptables规则
function delete_iptables_rule() {
    local chains=("INPUT" "FORWARD" "OUTPUT")
    local chain=$(dialog --menu "请选择要删除规则的链：" 10 40 3 "${chains[@]}" 3>&1 1>&2 2>&3 3>&-)
    local exit_status=$?
    if [ $exit_status -eq 1 ]; then
        dialog --msgbox "已取消操作或选择无效。" 6 40
        return
    fi

    local rule_number=$(dialog --inputbox "请输入要删除的规则编号 (如1、2、3...，按Enter跳过查看规则列表):" 8 40 3>&1 1>&2 2>&3 3>&-)
    if [ $? -eq 1 ]; then
        dialog --msgbox "已取消操作或输入无效。" 6 40
        return
    fi

    local rules=$(iptables -L $chain --line-numbers -n | grep '^ *[0-9]\+' | awk '{print $1, $NF}')
    local selected_rule=$(dialog --menu "请选择要删除的规则：" 20 60 10 $rules 3>&1 1>&2 2>&3 3>&-)
    if [ $? -eq 1 ]; then
        dialog --msgbox "已取消操作或选择无效。" 6 40
        return
    fi

    local rule_number_to_delete=$(echo "$selected_rule" | awk '{print $1}')
    local rule_to_delete=$(echo "$selected_rule" | awk '{$1=""; print $0}')

    if iptables -D $chain $rule_number_to_delete &>>"$LOG_FILE"; then
        log "INFO" "iptables规则已成功删除: $rule_to_delete"
        dialog --msgbox "iptables规则已成功删除。" 6 40
    else
        local error_code=$?
        log "ERROR" "删除iptables规则失败 (错误码: $error_code): $rule_to_delete"
        dialog --msgbox "删除iptables规则失败 (错误码: $error_code)，请检查输入并重试。" 6 40
    fi
}

# 主菜单
function main_menu() {
    while true; do
        local choice=$(dialog --menu "请选择操作类型：" 10 40 2 \
            1 "添加iptables规则" \
            2 "删除iptables规则" \
            3>&1 1>&2 2>&3 3>&-)

        case $choice in
            1) add_iptables_rule ;;
            2) delete_iptables_rule ;;
            *) dialog --msgbox "无效的选择。" 6 40 ;;
        esac
    done
}

# 启动主菜单
main_menu
